from collections.abc import Iterable
from typing import Callable, Generic, TypeVar

from zkl_aiutils_datasets.basics import Dataset
from zkl_aiutils_datasets.wrapping import wrap_dataset
from .chained import ChainedDatasets, ChainedIterator
from .delegating import DelegatingDataset, DelegatingIterator
from .mapped import MappedDataset, MappedIterator

SrcSample = TypeVar('SrcSample')
DstSample = TypeVar('DstSample')


class AugmentedIterator(DelegatingIterator[DstSample], Generic[DstSample, SrcSample]):
    def __init__(self,
        iterable: Iterable[SrcSample],
        augment_func: Callable[[SrcSample], Iterable[DstSample]],
    ):
        iterator = MappedIterator(iterable, augment_func)
        iterator = ChainedIterator(iterator)
        super().__init__(iterator)


class AugmentedDataset(DelegatingDataset[DstSample], Generic[DstSample, SrcSample]):
    def __init__(self,
        dataset: Dataset[SrcSample],
        augment_func: Callable[[SrcSample], Iterable[DstSample]],
    ):
        self._dataset = dataset
        self._augment_func = augment_func

        dataset = MappedDataset(dataset, augment_func)
        dataset = MappedDataset(dataset, wrap_dataset)
        dataset = ChainedDatasets(dataset)
        super().__init__(dataset)

    @property
    def dataset(self) -> Dataset[SrcSample]:
        return self._dataset

    @property
    def augment_func(self) -> Callable[[SrcSample], Iterable[DstSample]]:
        return self.augment_func
